import os
import numpy as np
import copy
import colorsys
from timeit import default_timer as timer
from keras import backend as K
from keras.models import load_model
from keras.layers import Input
from PIL import Image, ImageFont, ImageDraw
from nets.yolo3 import yolo_body, yolo_eval
from utils.utils import letterbox_image


# --------------------------------------------#
#   使用自己训练好的模型预测需要修改2个参数
#   model_path和classes_path都需要修改！
# --------------------------------------------#
class YOLO(object):
    my_model = 'ep096-loss11.385-val_loss14.261.h5'
    # my_model = 'yolo_weights.h5'
    if my_model == 'yolo_weights.h5':
        classes_path = 'coco_classes.txt'
    else:
        classes_path = 'voc_classes.txt'
    _defaults = {
        "model_path": 'demo/model_data/' + my_model,
        "anchors_path": 'demo/model_data/yolo_anchors.txt',
        "classes_path": 'demo/model_data/' + classes_path,
        "score": 0.5,
        "iou": 0.3,
        "max_boxes": 100,
        "model_image_size": (416, 416)
    }

    @classmethod
    def get_defaults(cls, n):
        if n in cls._defaults:
            return cls._defaults[n]
        else:
            return "Unrecognized attribute name '" + n + "'"

    # ---------------------------------------------------#
    #   初始化yolo
    # ---------------------------------------------------#
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)
        self.class_names = self._get_class()
        self.anchors = self._get_anchors()
        self.sess = K.get_session()
        self.boxes, self.scores, self.classes = self.generate()

    # ---------------------------------------------------#
    #   获得所有的分类
    # ---------------------------------------------------#
    def _get_class(self):
        classes_path = os.path.expanduser(self.classes_path)
        with open(classes_path) as f:
            class_names = f.readlines()
        class_names = [c.strip() for c in class_names]
        return class_names

    # ---------------------------------------------------#
    #   获得所有的先验框
    # ---------------------------------------------------#
    def _get_anchors(self):
        anchors_path = os.path.expanduser(self.anchors_path)
        with open(anchors_path) as f:
            anchors = f.readline()
        anchors = [float(x) for x in anchors.split(',')]
        return np.array(anchors).reshape(-1, 2)

    # ---------------------------------------------------#
    #   获得所有的分类
    # ---------------------------------------------------#
    global count
    count = 0
    def generate(self):
        global count
        model_path = os.path.expanduser(self.model_path)
        assert model_path.endswith('.h5'), 'Keras模型或权重必须是. h5文件'

        # 计算anchor数量
        num_anchors = len(self.anchors)
        num_classes = len(self.class_names)

        # 载入模型，如果原来的模型里已经包括了模型结构则直接载入。
        # 否则先构建模型再载入
        if count == 0:
            try:
                self.yolo_model = load_model(model_path, compile=False)
            except:
                self.yolo_model = yolo_body(Input(shape=(None, None, 3)), num_anchors // 3, num_classes)
                self.yolo_model.load_weights(self.model_path)
            else:
                assert self.yolo_model.layers[-1].output_shape[-1] == \
                       num_anchors / len(self.yolo_model.output) * (num_classes + 5), \
                    '模型与给定锚点和类大小不匹配'

            # print('{} model, anchors, and classes loaded.'.format(model_path))
            print('{} 模型、特征层锚和类别加载完毕。'.format(model_path))
            count = 1
        else:
            self.yolo_model = yolo_body(Input(shape=(None, None, 3)), num_anchors // 3, num_classes)
            self.yolo_model.load_weights(self.model_path)

        # 画框设置不同的颜色
        if model_path == 'model_data/yolo_weights.h5':
            hsv_tuples = [(x / len(self.class_names), 1., 1.)
                          for x in range(len(self.class_names))]
            self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
            self.colors = list(
                map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                    self.colors))

            # 打乱颜色
            np.random.seed(10101)
            np.random.shuffle(self.colors)
            np.random.seed(None)
        else:
            self.colors = [(255, 20, 255), (40, 100, 255), (0, 255, 255)]

        self.input_image_shape = K.placeholder(shape=(2,))

        boxes, scores, classes = yolo_eval(self.yolo_model.output, self.anchors,
                                           num_classes, self.input_image_shape, max_boxes=self.max_boxes,
                                           score_threshold=self.score, iou_threshold=self.iou)
        return boxes, scores, classes

    # ---------------------------------------------------#
    #   检测图片
    # ---------------------------------------------------#
    def detect_image(self, image):
        r_log2 = ''
        start = timer()

        # 调整图片使其符合输入要求
        new_image_size = (self.model_image_size[1], self.model_image_size[0])
        boxed_image = letterbox_image(image, new_image_size)
        image_data = np.array(boxed_image, dtype='float32')
        image_data /= 255.
        image_data = np.expand_dims(image_data, 0)  # Add batch dimension.

        # 预测结果
        out_boxes, out_scores, out_classes = self.sess.run(
            [self.boxes, self.scores, self.classes],
            feed_dict={
                self.yolo_model.input: image_data,
                self.input_image_shape: [image.size[1], image.size[0]],
                K.learning_phase(): 0
            })
        r_log1 = '在图片中找到 {} 处目标\n'.format(len(out_boxes))
        print('\033[0m\n在图片中找到 \033[34m{} \033[0m处目标\n'.format(len(out_boxes)))
        # 设置字体
        font = ImageFont.truetype(font='font/simhei.ttf',
                                  size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
        # 检测框厚度
        thickness = (image.size[0] + image.size[1]) // 600

        small_pic = []
        person, car, camera = 0, 0, 0
        for i, c in list(enumerate(out_classes)):
            predicted_class = self.class_names[c]
            box = out_boxes[i]
            score = out_scores[i]

            top, left, bottom, right = box
            top = top - 5
            left = left - 5
            bottom = bottom + 5
            right = right + 5

            top = max(0, np.floor(top + 0.5).astype('int32'))
            left = max(0, np.floor(left + 0.5).astype('int32'))
            bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
            right = min(image.size[0], np.floor(right + 0.5).astype('int32'))

            # 画框框
            label = '{} {:.2f}'.format(predicted_class, score)
            draw = ImageDraw.Draw(image)
            label_size = draw.textsize(label, font)
            label = label.encode('utf-8')
            labels = str(label)
            mydict = {'person': '\033[35m行人', 'car': '\033[34m机动车', 'camera': '\033[36m监控摄像头'}
            mydicts = {'person': '行人', 'car': '机动车', 'camera': '监控摄像头'}
            x = (out_boxes[i][1] + out_boxes[i][3]) / 2
            y = (out_boxes[i][0] + out_boxes[i][2]) / 2
            # print('x='+str(x),'y='+str(y))
            # print(out_boxes[i])
            width = image.size[0]
            height = image.size[1]
            # print('width=' + str(width), 'height=' + str(height))
            hx = 2.5
            y1 = width / hx
            y2 = width - y1
            x1 = height / hx
            x2 = height - x1
            # print('y1=' + str(y1), 'y2=' + str(y2), 'x1=' + str(x1), 'x2=' + str(x2))
            # print(y1,y2,x1,x2)
            # print(y,x1,x2)
            level = ''
            vertical = ''
            if x <= y1:
                level = '左'
            elif x >= y2:
                level = '右'
            if y <= x1:
                vertical = '上方'
            elif y >= x2:
                vertical = '下方'
            if level == '' and vertical == '':
                level = '中间区域'
            elif level != '' and vertical == '':
                level += '侧'
            res = level + vertical
            if labels[2:-6] == 'person':
                person += 1
            elif labels[2:-6] == 'car':
                car += 1
            elif labels[2:-6] == 'camera':
                camera += 1
            try:
                r_log2 += '第 {} 处有 {}% 的概率检测类别为：'.format(i + 1, round(score * 100, 2))+str(mydicts[labels[2:-6]])+'位于图片' + res + '<br>'
                print('\033[33m第 \033[34m{} \033[33m处有 \033[31m{}% \033[33m的概率检测类别为：'.format(i + 1, round(score * 100, 2)),mydict[labels[2:-6]], '\033[33m位于图片\033[31m' + res)
            except:
                r_log2 = '第 {} 处有 {}% 的概率检测类别为：'.format(i + 1, round(score * 100, 2))+str(labels[2:-6])+'位于图片' + res
                print('\033[33m第 \033[34m{} \033[33m处有 \033[31m{}% \033[33m的概率检测类别为：'.format(i + 1, round(score * 100, 2)),labels[2:-6], '\033[33m位于图片\033[31m' + res)

            if top - label_size[1] >= 0:
                text_origin = np.array([left, top - label_size[1]])
            else:
                text_origin = np.array([left, top + 1])

            for i in range(thickness):
                draw.rectangle(
                    [left + i, top + i, right - i, bottom - i],
                    outline=self.colors[c])
            draw.rectangle(
                [tuple(text_origin), tuple(text_origin + label_size)],
                fill=self.colors[c])
            draw.text(text_origin, str(label, 'UTF-8'), fill=(0, 0, 0), font=font)
            del draw

        end = timer()
        r_log3 = "检测到 行人 共 {} 个，机动车 共 {} 辆， 监控摄像头 共 {} 台".format(
                person, car, camera)
        print(
            "\n\033[33m检测到\033[35m 行人 \033[33m共\033[31m {} \033[33m个，\033[34m 机动车 \033[33m共\033[31m {} \033[33m辆，\033[36m 监控摄像头 \033[33m共\033[31m {} \033[33m台".format(
                person, car, camera))
        r_log4 = "预测用时:"+str(round(end - start, 3))+'秒'
        print("\033[33m预测用时:\033[34m", round(end - start, 3), '\033[33m秒\n')
        return image,r_log1,r_log2,r_log3,r_log4

    def close_session(self):
        self.sess.close()
